from config import *
from utils import *

insert_templates = [
    "Please observe the arrangement of all {category}s on the desktop. If the goal is to have all objects evenly spaced, where should an additional {category} be placed? Please mark the pixel location on the desktop.",
    "Examine the positioning of all {category}s on the desktop. To achieve equal spacing, where would you add another {category}? Indicate the pixel location on the desktop.",
    "Look at the current layout of all {category}s on the desktop. If we want them evenly distributed, where should an extra {category} be positioned? Please mark the pixel location.",
    "Consider the arrangement of all {category}s on the desktop. To ensure uniform spacing, where should an additional {category} be added? Indicate the pixel location.",
    "Review the placement of all {category}s on the desktop. If the aim is equal spacing, where would you place another {category}? Please mark the pixel location.",
    "Observe how the {category}s are arranged on the desktop. To make them evenly spaced, where should an extra {category} be located? Indicate the pixel position.",
    "Analyze the distribution of all {category}s on the desktop. If we want equal spacing, where should an additional {category} be placed? Please mark the pixel location.",
    "Inspect the current setup of {category}s on the desktop. To achieve uniform spacing, where would you add another {category}? Indicate the pixel location.",
    "Take a look at the arrangement of all {category}s on the desktop. If the goal is even spacing, where should an extra {category} be positioned? Please mark the pixel location.",
    "Evaluate the positioning of {category}s on the desktop. To ensure equal spacing, where would you place another {category}? Indicate the pixel location.",
    "Please assess the layout of all {category}s on the desktop. If we aim for uniform spacing, where should an additional {category} be added? Mark the pixel location.",
    "Examine the current configuration of {category}s on the desktop. To achieve even spacing, where would you insert another {category}? Indicate the pixel position.",
    "Look into the arrangement of {category}s on the desktop. If equal spacing is desired, where should an extra {category} be placed? Please mark the pixel location.",
    "Consider the current distribution of {category}s on the desktop. To make them evenly spaced, where would you add another {category}? Indicate the pixel location.",
    "Review how the {category}s are positioned on the desktop. If the aim is uniform spacing, where should an additional {category} be located? Mark the pixel position.",
    "Observe the setup of all {category}s on the desktop. To ensure equal spacing, where would you place another {category}? Indicate the pixel location.",
    "Analyze the current layout of {category}s on the desktop. If we want them evenly spaced, where should an extra {category} be added? Please mark the pixel location.",
    "Inspect the arrangement of {category}s on the desktop. To achieve uniform spacing, where would you insert another {category}? Indicate the pixel position.",
    "Take note of how the {category}s are arranged on the desktop. If equal spacing is the goal, where should an additional {category} be placed? Mark the pixel location.",
    "Evaluate the current positioning of {category}s on the desktop. To make them evenly spaced, where would you add another {category}? Indicate the pixel location."
]

side_templates = [
    "In the current image, all {category}s are evenly spaced. To maintain this equal spacing, where should an additional {category} be placed on the {side}? Please mark the pixel location on the desktop.",
    "Observe the arrangement of all {category}s in the image. To keep them evenly spaced, where would you place another {category} on the {side}? Indicate the pixel location on the desktop.",
    "Given the evenly spaced {category}s in the image, where should an extra {category} be added on the {side} to preserve the spacing? Please mark the pixel location.",
    "All {category}s are currently evenly spaced. To continue this pattern, where would you add another {category} on the {side}? Indicate the pixel location on the desktop.",
    "Considering the uniform spacing of {category}s, where should an additional {category} be positioned on the {side}? Please mark the pixel location.",
    "The {category}s are evenly distributed in the image. To maintain this, where would you place another {category} on the {side}? Indicate the pixel location on the desktop.",
    "With all {category}s evenly spaced, where should an extra {category} be placed on the {side} to continue the pattern? Please mark the pixel location.",
    "Looking at the current layout of {category}s, where would you add another {category} on the {side} to keep the spacing even? Indicate the pixel location on the desktop.",
    "To preserve the equal spacing of {category}s, where should an additional {category} be positioned on the {side}? Please mark the pixel location.",
    "Given the current even spacing of {category}s, where would you place another {category} on the {side}? Indicate the pixel location on the desktop.",
    "All {category}s are arranged with equal spacing. To maintain this, where should an extra {category} be added on the {side}? Please mark the pixel location.",
    "Considering the evenly spaced {category}s, where would you position another {category} on the {side} to continue the arrangement? Indicate the pixel location on the desktop.",
    "To keep the {category}s evenly spaced, where should an additional {category} be placed on the {side}? Please mark the pixel location.",
    "With the current equal spacing of {category}s, where would you add another {category} on the {side}? Indicate the pixel location on the desktop.",
    "Observing the uniform arrangement of {category}s, where should an extra {category} be positioned on the {side}? Please mark the pixel location.",
    "Given the evenly spaced {category}s, where would you place another {category} on the {side} to maintain the pattern? Indicate the pixel location on the desktop.",
    "All {category}s are currently evenly spaced. To continue this arrangement, where should an additional {category} be added on the {side}? Please mark the pixel location.",
    "Considering the equal spacing of {category}s, where would you position another {category} on the {side}? Indicate the pixel location on the desktop.",
    "To preserve the uniform spacing of {category}s, where should an extra {category} be placed on the {side}? Please mark the pixel location.",
    "With the current layout of evenly spaced {category}s, where would you add another {category} on the {side}? Indicate the pixel location on the desktop."
]

def find_missing_number(arr):
    sorted_arr = sorted(arr)
    for i in range(len(sorted_arr) - 1):
        if sorted_arr[i + 1] - sorted_arr[i] > 1:
            return sorted_arr[i] + 1
    return None

def gen_extra1_qa(obj_dict, cam, table):
    qa = []
    arr = []
    ID = None
    for id, info in obj_dict.items():
        if '_' in id and info['visible']:
            arr.append(eval(id.split('_')[1]))
            ID = id.split('_')[0]
    if len(arr) != max(arr) - min(arr) + 1:
        x = find_missing_number(arr)
        left = obj_dict[f'{ID}_{x-1}']
        right = obj_dict[f'{ID}_{x+1}']
        cate = left['minimal_desc']
        left_loc = left['3dtabloc']
        right_loc = right['3dtabloc']
        mid_loc = ((left_loc[0] + right_loc[0]) / 2, (left_loc[1] + right_loc[1]) / 2)
        pix = world_to_camera_pixel((mid_loc[0], mid_loc[1], table[4]), cam)
        question = choice(insert_templates).format(category = cate)
        if random() < 0.4:
            question = question + SUFFIX
        answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
        thinking = f'All the {cate}s are arranged in a sequence at:'
        for i in arr:
            thinking += f'{format_coordinate(obj_dict[f"{ID}_{i}"])},'
        thinking += f'The {cate} at the missing position should be placed at {format_coordinatexy(pix[0], pix[1])}.'
        qa = [{
            'question': question,
            'answer': answer,
            'thinking': thinking,
        }]
    else:
        arr = sorted(arr)
        obj1 = obj_dict[f'{ID}_{arr[0]}']
        obj2 = obj_dict[f'{ID}_{arr[1]}']
        # calc left loc, let (left loc, obj1, obj2) be a line
        loc1 = obj1['3dtabloc']
        loc2 = obj2['3dtabloc']
        left_loc = (loc1[0] - (loc2[0] - loc1[0]), loc1[1] - (loc2[1] - loc1[1]))
        pix = check_table_pos((left_loc[0], left_loc[1], table[4]), cam, table)
        if pix != None:
            question = choice(side_templates).format(category = obj1['minimal_desc'], side = 'left')
            if random() < 0.4:
                question = question + SUFFIX
            answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
            thinking = f'All the {obj1["minimal_desc"]}s are arranged in a sequence at:'
            for i in arr:
                thinking += f'{format_coordinate(obj_dict[f"{ID}_{i}"])},'
            thinking += f'The {obj1["minimal_desc"]} at the left side should be placed at {format_coordinatexy(pix[0], pix[1])}.'
            qa.append({
                'question': question,
                'answer': answer,
                'thinking': thinking,
            })
        obj1 = obj_dict[f'{ID}_{arr[-2]}']
        obj2 = obj_dict[f'{ID}_{arr[-1]}']
        # calc right loc, let (obj1, obj2, right loc) be a line
        loc1 = obj1['3dtabloc']
        loc2 = obj2['3dtabloc']
        right_loc = (loc2[0] + (loc2[0] - loc1[0]), loc2[1] + (loc2[1] - loc1[1]))
        pix = check_table_pos((right_loc[0], right_loc[1], table[4]), cam, table)
        if pix != None:
            question = choice(side_templates).format(category = obj2['minimal_desc'], side = 'right')
            if random() < 0.4:
                question = question + SUFFIX
            answer = f'[{format_coordinatexy(pix[0], pix[1])}]'
            thinking = f'All the {obj2["minimal_desc"]}s are arranged in a sequence at:'
            for i in arr:
                thinking += f'{format_coordinate(obj_dict[f"{ID}_{i}"])},'
            thinking += f'The {obj2["minimal_desc"]} at the right side should be placed at {format_coordinatexy(pix[0], pix[1])}.'
            qa.append({
                'question': question,
                'answer': answer,
                'thinking': thinking,
            })
    return qa

    